import os
import argparse
import threading
import queue
from flask import Flask, request, jsonify
from threading import Thread
from .scheduler import TaskScheduler
from .config import Config
from .models import DLTask


app = Flask(__name__)
config = Config.load()
scheduler = TaskScheduler(config)
scheduler_thread = Thread(target=scheduler.scheduling_loop, daemon=True)
scheduler_thread.start()

@app.route('/health', methods=['GET'])
def health_check():
    return '', 200  # 返回空内容和200状态码表示服务正常

@app.route('/add-task', methods=['POST'])
def add_task():
    try:
        data = request.json
        command = data.get('command')
        device = data.get('device', [])
        num_gpu = data.get('num_gpu', 0)
        mem_cpu = data.get('mem_cpu')
        mem_per_gpu = data.get('mem_per_gpu')
        max_retries = data.get('max_retries', config.max_retries)
        env = data.get('env', os.environ.copy())
        pwd = data.get('pwd', "~/")

        # if num_gpu > 0 and len(device) == 0:
        #     return jsonify({"error": "当请求GPU时，必须指定具体的设备"}), 400
        
        for dev in device:
            if dev not in config.available_devices:
                return jsonify({"error": f"设备 {dev} 不在可用设备列表中"}), 400
        
        task = DLTask(
            command=command,
            num_gpu=num_gpu,
            mem_cpu=mem_cpu,
            mem_per_gpu=mem_per_gpu,
            require_devices=device if device else None,
            max_retries=max_retries,
            env=env,
            pwd=pwd,
        )
        scheduler.add_task(task)
        return jsonify({"message": f"任务已添加 (ID: {task.task_id[:8]})"}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/status', methods=['GET'])
def status():
    try:
        gpu_table_data = []
        for device in config.available_devices:
            used = scheduler.gpu_allocated_mem.get(device, 0) / 1024
            total = scheduler.gpu_total_mem.get(device, 0) / 1024
            gpu_table_data.append({
                "device": device,
                "used": f"{used:.1f}",
                "total": f"{total:.1f}"
            })

        cpu_used = scheduler.cpu_allocated_mem / 1024 # GB
        cpu_reserved = config.reserved_cpu_mem / 1024 # GB
        cpu_available = config.cpu_total_mem - cpu_reserved
        cpu_utilized = min(cpu_used, cpu_available)

        task_data = []
        with scheduler._lock:
            for task in scheduler.running_tasks + list(scheduler.pending_tasks) + scheduler.finished_tasks:
                task_data.append({
                    "id": task.task_id[:8],
                    "status": task.status,
                    "command": task.command[:40] + ("..." if len(task.command)>40 else ""),
                    "gpus": ",".join(task.allocated_devices) if task.allocated_devices else "-",
                    "cpu-mem": task.mem_cpu,
                    "retries": f"{task.retries}/{task.max_retries}"
                })

        return jsonify({
            "gpu_status": gpu_table_data,
            "cpu_status": {
                "total_mem": config.cpu_total_mem,
                "reserved_mem": cpu_reserved,
                "available_mem": cpu_available,
                "used_mem": cpu_utilized
            },
            "tasks": task_data
        }), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500



@app.route('/logs/<task_id>', methods=['GET'])
def logs(task_id):
    try:
        matched_logs = []
        for filename in os.listdir(config.log_dir):
            if filename.endswith(".log") and task_id.lower() in filename.lower():
                matched_logs.append(filename)
                
        if not matched_logs:
            return jsonify({"error": f"未找到匹配 '{task_id}' 的任务日志"}), 404
        elif len(matched_logs) > 1:
            return jsonify({
                "error": "找到多个匹配的任务日志，请指定更精确的任务ID",
                "matched_task_logs": matched_logs
            }), 400
        else:
            # 唯一匹配文件时的处理
            full_path = os.path.join(config.log_dir, matched_logs[0])
            result_queue = queue.Queue()
            error_queue = queue.Queue()
            
            def read_file():
                try:
                    with open(full_path, "r") as f:  # 使用 with 确保文件句柄释放
                        content = f.read()
                        result_queue.put(content)
                except Exception as e:
                    error_queue.put(str(e))
            
            read_thread = threading.Thread(target=read_file)
            read_thread.start()
            read_thread.join(timeout=5)  # 设置 5 秒超时
            
            if read_thread.is_alive():
                return jsonify({"error": "日志文件被锁定，暂时无法读取"}), 423  # 使用 423 Locked 状态码
            elif not error_queue.empty():
                return jsonify({"error": error_queue.get()}), 400
            else:
                return jsonify({
                    "log_content": result_queue.get(),
                    "file_name": matched_logs[0]
                }), 200
                
    except FileNotFoundError:
        return jsonify({"error": f"日志目录不存在"}), 404
    except Exception as e:
        return jsonify({"error": f"服务器内部错误: {str(e)}"}), 500

@app.route('/config-update', methods=['POST'])
def config_update():
    try:
        data = request.json
        available_devices = data.get('available_devices')
        reserved_mem = data.get('reserved_mem')
        reserved_cpu_mem = data.get('reserved_cpu_mem')
        max_retries = data.get('max_retries')

        if available_devices:
            new_devices = available_devices.split(",")
            config.available_devices = new_devices
            scheduler._init_gpu_metadata(config.available_devices)
        
        if reserved_mem is not None:
            if reserved_mem < 0:
                return jsonify({"error": "预留内存不能为负数"}), 400
            config.reserved_mem = reserved_mem
        
        if reserved_cpu_mem is not None:
            if reserved_cpu_mem < 0 or reserved_cpu_mem > config.cpu_total_mem:
                return jsonify({"error": f"预留CPU内存应在0到{config.cpu_total_mem}之间"}), 400
            config.reserved_cpu_mem = reserved_cpu_mem
        
        if max_retries is not None:
            if max_retries < 0:
                return jsonify({"error": "最大重试次数不能为负数"}), 400
            config.max_retries = max_retries
        
        config.save()
        return jsonify({"message": "配置已更新", "config": config.model_dump_json(indent=2)}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/stop', methods=['POST'])
def stop_scheduler():
    try:
        scheduler.stop()
         # 定义一个函数来发送停止信号
        def send_kill_signal():
            # 延迟一段时间发送信号，确保响应已经返回
            import signal
            import time
            time.sleep(1)  # 延迟1秒
            os.kill(os.getpid(), signal.SIGINT)
        
        # 启动一个新线程来发送信号
        thread = Thread(target=send_kill_signal)
        thread.start()
        return jsonify({"message": "调度器已停止"}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

def run(*args, **kwargs):
    # 使用解析后的参数启动Flask应用
    app.run(*args, **kwargs)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Run the Flask application with custom host and port.")
    parser.add_argument("--host", default="127.0.0.1", help="Host address to bind the server to")
    parser.add_argument("--port", type=int, default=12030, help="Port number to bind the server to")

    args = parser.parse_args()

    # 使用解析后的参数启动Flask应用
    app.run(host=args.host, port=args.port)